#!/usr/bin/env python

# -------
# Imports
# -------
import config

# ROS 
import roslib; roslib.load_manifest(config.NODE_NAME)
import rospy
from geometric_image_filter.msg import FilteredImage
from sensor_msgs.msg import Image, CameraInfo, LaserScan, PointCloud, RegionOfInterest
from geometry_msgs.msg import Point, PointStamped
import tf

# General  python
import threading
import sys
from numpy import *
from collections import deque

# ------------------------------------------
class GeometricImageFilter(threading.Thread):
# ------------------------------------------
	# -------		
	# Members:
	# -------
	filteredImagePub = None 		# Publisher for the FilteredImage messages this node outputs
	transformListener = None		# This is used to access the tf frames

	# Stored messages & relevant information
	prevCameraInfoMsg = None 		# The last camera info message that we
	cameraInfoMsg = None			# The most recent camera info message we've received 
	imageMsg = None					# The most recent image message we've received
	stereoPointCloudMsg = None		# The most recent stereo data we've received
	laserScanMsg = None				# The most recent laser data we've received (laser scan)
	cameraPtWorld = None

	# Filtering variables
	windowSizes = None								# A list of the object window sizes
	regionsOfInterest = None						# A collection of *all* possible regions of interest for this camera
	groundNormalWorld = array([[0],[0],[1],[0]])	# The ground plane normal vector (in the world coordinate frame)
	dS = None
	dH = None
	dC = None 
	ry = None
	
	# ----------------
	def __init__(self):
	# ----------------
		# Super
		threading.Thread.__init__(self)

		# Initialize the node
		rospy.init_node(config.NODE_NAME)

		# Subscriptions
		rospy.Subscriber(config.CAMERA_INFO_TOPIC, CameraInfo, self.__cameraInfoCallback)
		rospy.Subscriber(config.IMAGE_TOPIC, Image, self.__imageCallback)
		rospy.Subscriber(config.STEREO_PT_CLOUD_TOPIC, PointCloud, self.__stereoPointCloudCallback)
		rospy.Subscriber(config.LASER_SCAN_TOPIC, LaserScan, self.__laserScanCallback)
		self.transformListener = tf.TransformListener()

		# We're going to publish filtered images
		self.filteredImagePub = rospy.Publisher(config.FILTERED_IMAGE_TOPIC, FilteredImage)

	# ----
	# Core|
	# -----------
	def run(self):
	# -----------
		""" The main control loop for the geometric_image_processing node. """

		from time import time

		# Status info
		rospy.loginfo('Starting %s node' % (config.NODE_NAME))

		# Preprocessing
		# -------------
		# Block for a camera info message
		rospy.loginfo('Blocking for camera calibration information...')
		while self.cameraInfoMsg == None: pass
		cameraInfoMsg = self.cameraInfoMsg

		# Determine the minimum object size in *pixels*
		minWindowSize = self.__calculateMinWindowSize(cameraInfoMsg.K, config.MIN_SIZE_M, config.MAX_DEPTH_M)
		rospy.loginfo('Minimum window size: %ix%i pixels', minWindowSize, minWindowSize)

		# Precompute the set of all possible regions of interest
		rospy.loginfo('Generating set of possible regions of interest')
		self.windowSizes = [v for v in self.__windowSizeRange(minWindowSize, min(cameraInfoMsg.width, cameraInfoMsg.height), config.WINDOW_SCALE_FACTOR)]
		self.regionsOfInterest = self.__generateRegionsOfInterest(self.windowSizes, self.cameraInfoMsg.width, self.cameraInfoMsg.height)

		# Main control loop
		# -----------------
		times = []
		averages = []
		rospy.loginfo('Entering main control loop')
		while not rospy.is_shutdown():

			# Collect Messages:
			# ----------------
			imageMsg, cameraInfoMsg, cameraPtWorld, groundNormalCamera, stereoPointCloudMsg, laserScanMsg = self.__getMessages()
	
			# Make sure we have enough info to proceed
			if imageMsg != None and cameraInfoMsg != None and cameraPtWorld != None:
				###################
				startTime = time() 
				###################

				updateD = False	
		
				# Update size constraint:	
				# ----------------------
				# If necessary, compute the depths consistent with the possible image window sizes and the min/max feasible object size
				if self.prevCameraInfoMsg == None or self.prevCameraInfoMsg.K != cameraInfoMsg.K:
					rospy.loginfo('Computing the range of feasible depths given image window size and feasible object size')
					self.dS = self.computeSizeConstraint(cameraInfoMsg, self.windowSizes, config.MIN_SIZE_M, config.MAX_SIZE_M)
					self.dS = self.__filterDepthConstraints(self.dS, config.MAX_DEPTH_M)

					# Update the projections of v into the real world 	
					self.ry = self.__computeProjection(cameraInfoMsg)					
	
					# Update the previously processed camera info message
					self.prevCameraInfoMsg = cameraInfoMsg

					# Since we've changed dS, we're going to have to update dC
					updateD = True

				# Update height constraint:
				# ------------------------
				# If necessary, compute the depths consistent with the possible image heights and the min/max real world feasible object height
				if self.cameraPtWorld == None or cameraPtWorld[2] != self.cameraPtWorld[2]:
					rospy.loginfo('Computing the range of feasible depths given image height and feasible object height')
					self.dH = self.computeHeightConstraint(cameraInfoMsg, groundNormalCamera, self.ry, cameraPtWorld[2], config.MIN_HEIGHT_M, config.MAX_HEIGHT_M)
					self.dH = self.__filterDepthConstraints(self.dH, config.MAX_DEPTH_M)

					# Update the camera point
					self.cameraPtWorld = cameraPtWorld
				
					# Since we've changed dH, we're going to have to update dC 
					updateD = True

				# Update the combined depth constraints:
				# -------------------------------------
				if updateD:
					rospy.loginfo('Computing the intersection of the size & height constraints')					
					self.dC = self.combineDepthConstraints(self.dS, self.dH)

				# Filter the set of possible ROI's:
				# --------------------------------
				filteredROIS = self.filterRegionsOfInterest(self.regionsOfInterest, self.dC, None, None)
				rospy.loginfo('Number of filtered regions of interest: %i', len(filteredROIS))

				# Create our FilteredImage message
				filteredImage = FilteredImage(image = imageMsg, rois = filteredROIS)
				self.filteredImagePub.publish(filteredImage)

				# Timing information
				duration = time() - startTime
				rospy.loginfo('Duration: %s', str(duration))
				times.append(duration)
				averages.append(float(len(filteredROIS)) / self.totalROI)

		rospy.loginfo('Shutting down %s node' % (config.NODE_NAME))

		rospy.loginfo('Average duration: %s', (sum(times) / len(times))) 
		rospy.loginfo('Average percent of total: %f', sum(averages) / len(averages))
		rospy.loginfo('Total number: %s', len(times))

	# -------------------------------------------------------------------------------------
	def computeSizeConstraint(self, cameraInfoMsg, windowSizes, minWorldSize, maxWorldSize):
	# -------------------------------------------------------------------------------------
		# Get the focal length, fx
		fx = cameraInfoMsg.K[0]		

		# Common values
		sMinFx = fx * minWorldSize
		sMaxFx = fx * maxWorldSize
		
		# Compute the range of feasible depths
		feasibleSizes = {}
		for w in windowSizes:
			val = ((sMinFx / w), (sMaxFx / w))
			feasibleSizes[w] = val		

		return feasibleSizes

	# -------------------------------------------------------------------------------------------------------------
	def computeHeightConstraint(self, cameraInfoMsg, groundNormalCamera, ry, cameraHeightM, minHeightM, maxHeightM):
	# -------------------------------------------------------------------------------------------------------------
		# Common values
		dH = {}
		hMinDiff = -1 * (cameraHeightM - minHeightM)
		hMaxDiff = -1 * (cameraHeightM - maxHeightM)

		# Alias the ground plane's normal vector
		n = groundNormalCamera

		# Comptute the depth constraint
		for v in range(0, cameraInfoMsg.height):
			nr = (float(n[1]) * ry[v]) + float(n[2])
			val = ((hMinDiff / nr), (hMaxDiff / nr)) 

			dH[v] = val

		return dH

	# ---------------------------------------
	def combineDepthConstraints(self, dS, dH):
	# ---------------------------------------
		""" Combines the size and height geometric constraints into a single collection.
		Specifically, this method combines dS and dH in to a mapping of (w,v) --> (min, max depth).
		"""

		# The combination of dS and dH
		dC = {}

		# Combine the constraints
		for w, sizeRange in dS.iteritems():
			for v, heightRange in dH.iteritems():
				dC[w,v] = self.__rangeIntersection(sizeRange, heightRange)

		return dC

	# -------------------------------------------------
	def filterRegionsOfInterest(self, rois, dC, mS, mL):
	# -------------------------------------------------
		# The filtered collection of rois
		filtered = []

		for key, value in rois.iteritems():
			if dC[key] != None:
				filtered += value # NOTE TO SELF: Eliminating this line makes all the difference in the world.
		
		return filtered

	# --------
	# Helpers |
	# ---------------------
	def __getMessages(self):	
	# ---------------------
		""" A private function that compiles a snapshot of all of the relevant data needed to run the geometric filter. """

		# The variables we're going to try and fill
		imageMsg = None
		stereoPointCloudMsg = None
		cameraInfoMsg = None
		laserScanMsg = None
		cameraPtWorld = None
		groundNormalCamera = None

		# Make sure we have an image
		if self.imageMsg != None:
			imageMsg = self.imageMsg

			# Make sure we have a camera info message and that it's current
			ts = lambda ts: ts.secs + (ts.nsecs / 1000000000)
			imageTs = ts(imageMsg.header.stamp)
			if self.cameraInfoMsg != None and abs(imageTs - ts(self.cameraInfoMsg.header.stamp)) < config.IMG_CAMINFO_TOLERANCE_S:
				cameraInfoMsg = self.cameraInfoMsg

				# Make sure that we have a camera point					
				cameraPtWorld, groundNormalCamera = self.__getCameraOrientation(cameraInfoMsg.header.frame_id, config.BASE_FRAME)
				if cameraPtWorld != None:
					
					# If we have stereo data, include it
					if self.stereoPointCloudMsg != None and abs(imageTs - ts(self.stereoPointCloudMsg.header.stamp)) < config.IMG_STEREO_TOLERANCE_S:
						stereoPointCloudMsg = self.stereoPointCloudMsg

					# Include laser data if available
					if self.laserScanMsg != None and abs(imageTs - ts(self.laserScanMsg.header.stamp)) < config.IMG_LASER_TOLERANCE_S:
						laserScanMsg = self.laserScanMsg	

		# Return the messages we've compiled
		return (imageMsg, cameraInfoMsg, cameraPtWorld, groundNormalCamera, stereoPointCloudMsg, laserScanMsg)

	# ------------------------------------------------------------
	def __filterDepthConstraints(self, depthConstraints, maxDepth):
	# ------------------------------------------------------------
		""" Filters a list of depth constraints in the form (min, max depth).
		Depth constraints are filtered out under one of two conditions:
		1.) The minimum depth is greater than the maximum allowed depth
		2.) The maximum depth is negative
		"""

		for i, v in depthConstraints.iteritems():
			if v[1] < 0 or v[0] > maxDepth:
				depthConstraints[i] = None
		return depthConstraints	

	# ------------------------------------------
	def __computeProjection(self, cameraInfoMsg):
	# ----------------------------------------- 
		""" This helper method computes the y component of the real-world ray projection for every value of v. """

		# Reshape and invert k
		k = self.__reshapeKMatrix(cameraInfoMsg.K)
		kPrime = linalg.inv(k)

		return [((kPrime[1,1]*v + kPrime[1,2])/(kPrime[2,1]*v + kPrime[2,2])) for v in range(0, cameraInfoMsg.height)]						

	# ------------------------------------------------------
	def __getCameraOrientation(self, cameraFrame, baseFrame):
	# ------------------------------------------------------
		""" A private function that retrieves the position of the camera relative to some base frame.
		@param cameraFrame The tf frame of the camera.
		@param baseFrame The tf frame of the base.	

		@return a 2-tuple -- a numpy array indicating the camera's origin (x,y,z,1) and a numpy array indicating the normal vector of the ground plane (x,y,z,1)
		"""

		try:
			# Make sure our two frames exist
			if self.transformListener.frameExists(cameraFrame) and self.transformListener.frameExists(baseFrame):
				# Get the transformation from camera -> base frame (cb) 
				latestCommonTime = self.transformListener.getLatestCommonTime(cameraFrame, baseFrame)
				
				# Compute the camera's origin point in the base frame	
				cameraOriginBase, cbRotation = self.transformListener.lookupTransform(baseFrame, cameraFrame, latestCommonTime)

				# Compute the ground plane's normal vector in the camera's frame 
				bcTranslation, bcRotation = self.transformListener.lookupTransform(cameraFrame, baseFrame, latestCommonTime)
				bcRotationMatrix = self.transformListener.fromTranslationRotation((0,0,0), bcRotation)
				groundNormalCamera = dot(bcRotationMatrix, self.groundNormalWorld) 

				# -------------------------------------------------------------------------------------------------------------------
				# TODO: Should probably check here to make sure that groundNormalCamera[0] (x) ~ 0, otherwise the camera has 'rolled'
				# -------------------------------------------------------------------------------------------------------------------

				return (cameraOriginBase, groundNormalCamera)
			
			# Otherwise return None indicating we can't get the camera point
			else:
				return (None, None)

		# If anything goes wrong return None indicating we can't get the camera point
		except tf.Exception as e:
			return (None, None)

	# --------------------------------------------------------
	def __windowSizeRange(self, minSize, maxSize, scaleFactor):
	# --------------------------------------------------------
		""" Generates window sizes given a min/max size and a scale factor.
		Values are generated according to the formula size(i) = minSize(scaleFactor)^(i-1). 
		"""
		val = minSize
		yield val 
		while val < maxSize:
			val *= scaleFactor
			if val <= maxSize: yield int(math.floor(val))

	# ---------------------------------------------------------------------
	def __generateRegionsOfInterest(self, windowSizes, imgWidth, imgHeight):
	# ---------------------------------------------------------------------
		""" Generates a collection of all possible regions of interest. """

		# A mapping of (w,v) --> [RegionOfInterest, RegionOfInterest, ...] 
		regionsOfInterest = {}

		total = 0

		for w in windowSizes:
			# Precompute this commonly used value		
			halfWidth = int(w/2)

			for v in range(halfWidth, imgHeight - halfWidth):
				yOffset = v - halfWidth
				regionsOfInterest[w,v] = [RegionOfInterest(x_offset = u - halfWidth, y_offset = yOffset, height = w, width = w) for u in range(halfWidth, imgWidth - halfWidth)]
				#regionsOfInterest[w,v] = [(u - halfWidth, yOffset,w) for u in range(halfWidth, imgWidth - halfWidth)]
				
				total += len(regionsOfInterest[w,v])

		################################################################
		rospy.loginfo('Total number of regions of interest: %i', total)
		self.totalROI = total
		################################################################

		return regionsOfInterest

	# ---------------------------
	def __reshapeKMatrix(self, k):
	# ---------------------------
		""" Simple helper function to reshape a flat K matrix (presumably from a CameraInfo message) into a 3x3 numpy array. """
		if len(k) == 9:
			return array([k[0:3], k[3:6], k[6:9]])
		else: 
			raise Exception('Unable to reshape K matrix: %s' % str(k))

	# ----------------------------------------------------------------------
	def __calculateMinWindowSize(self, k, minAllowedSizeM, maxAllowedDepthM):
	# ----------------------------------------------------------------------
		""" Calculate the minimum object window size (in pixels). 
		This calculation is achieved by projecting two real world points:
		(0, 0, maxAllowedDepth) and (minAllowedSizeM, 0, maxAllowedDepth) from 3D 
		to 2D and then calculating the distance between them on the x axis.
		"""

		# Reshape and invert k
		k = self.__reshapeKMatrix(k)

		# Create our two points
		pt3A = array([[0], [0], [maxAllowedDepthM]]) 
		pt3B = array([[minAllowedSizeM], [0], [maxAllowedDepthM]]) 

		# Project the points into image space
		pt2A = dot(k, pt3A)
		pt2B = dot(k, pt3B)
	
		# Calculate the distance 
		return int(abs(pt2A[0] - pt2B[0]))

	# -----------------------------------
	def __rangeIntersection(self, r1, r2):
	# -----------------------------------
		""" Computer the intersection of two one dimensional ranges. """
		if r1 and r2:
			r = (max(r1[0], r2[0]), min(r1[1], r2[1]))
			return r if r[1] > r[0] else None 
		else: return None

	# ----------
	# Callbacks |
	# ---------------------------------
	def __cameraInfoCallback(self, msg):
	# ---------------------------------
		# Store the most recent camera info message we've received
		self.cameraInfoMsg = msg

	# ----------------------------
	def __imageCallback(self, msg):
	# ----------------------------
		# Store the most recent image we've received
		self.imageMsg = msg

	# ---------------------------------------
	def __stereoPointCloudCallback(self, msg):
	# ---------------------------------------
		# Store the most recent stereo point cloud we've received
		self.stereoPointCloudMsg = msg

	# --------------------------------
	def __laserScanCallback(self, msg):
	# --------------------------------
		# Store the most recent laser data (laser scan) that we've received
		self.laserScanMsg = msg

# -----------------------
if __name__ == '__main__':
# -----------------------
	controller = GeometricImageFilter()
	controller.run()
